import torch
from torch import nn
from d2l import torch as d2l
import sys
sys.path.append('E:\AI\DL')
from Train_ import train_func as tf
from SoftmaxRegression import data_generate as dg

net = nn.Sequential(nn.Flatten(),nn.Linear(784,256),nn.ReLU(),nn.Linear(256,10))
def init_weights(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight,0.01)
net.apply(init_weights)

num_workers = 4
batch_size,lr,num_epochs = 256,0.1,10
loss_func = nn.CrossEntropyLoss()
trainer = torch.optim.SGD(net.parameters(),lr=lr)
train_iter,test_iter = dg.load_data_fashion_mnist(batch_size,num_workers)
tf.train(net,train_iter,test_iter,loss_func,num_epochs,trainer)

